#include <string>
#include <iostream>
#include "log.h"
#include "netpcap.h"
#include "netpcap_intf.h"
#include "ip_packet.h"
#include "l4_packet.h"
#include "tcp_packet.h"
#include "tcpip_init.h"
#include "session.h"
#include <fstream>
#include "protocol_detector.h"
#include "protocol_log.h"

using namespace secsmart_pcap;
using namespace secsmart_tcpip;
using namespace secsmart_session;

#define MYSQL_MATCH_THRESHOLD   5


uint32_t g_maxThreshold = MYSQL_MATCH_THRESHOLD;
std::string g_target_file;
//bool g_withlog = false ;

extern void sniffer_db_app_print_version();

// tcp 连接建立事件处理
std::shared_ptr<void> tcp_conn_event_start_proc(const tcp_stream_sptr conn, void *user_data)
{
    (void)user_data;
    session_sptr s = std::make_shared<session_t>();
    s->sip = conn->get_sip();
    s->sport = conn->get_sport();
    s->dip = conn->get_dip();
    s->dport = conn->get_dport();
    s->matchCount = 0;
    s->unknowPktCount = 0;
    s->dbVersion = "UNKNOWN";
    s->detector = std::make_shared<secsmart_protocol::ProtocolDetector>(g_target_file, g_maxThreshold);
    if (conn->rcv_syn) {
        s->direction = SESSION_DIRCTION_REQUEST;
        LOG_DEBUG("create conn source {}:{} --> destination {}:{}", conn->get_sip().toString(), s->sport, conn->get_dip().toString(), s->dport);
    } else {
        s->direction = SESSION_DIRCTION_UNKNOW;
        LOG_DEBUG("create conn {}:{} -- {}:{}", conn->get_sip().toString(), s->sport, conn->get_dip().toString(), s->dport);
    }
    
    return s;
}


// tcp 连接结束事件处理
void tcp_conn_event_end_proc(const tcp_stream_sptr conn, uint32_t close_reason, std::shared_ptr<void> usr_conn_data,
                               void *usr_data)
{
    (void)usr_data;
    (void)usr_conn_data;
    uint32_t rcv_pkt_num = conn->statistic.pkt_rcv_num;

    std::string conn_info;
    if (conn->is_in_tunnel()) {
        // [] 中为外层ip地址
        LOG_INFO("tcp conn end reason:{}. source {}[{}]:{} --> destination {}[{}]:{}. rcv_pkt_num:{}", close_reason,
                 conn->get_sip().toString(), conn->get_outsip().toString(), conn->get_sport(),
                 conn->get_dip().toString(), conn->get_outdip().toString(), conn->get_dport(), rcv_pkt_num);

        conn_info = "{ \"tcp_conn_end\" : \"" + conn->get_sip().toString() + "[" + conn->get_outsip().toString() + "]" +
                    ":" + std::to_string(conn->get_sport()) + "-->" + conn->get_dip().toString() + "[" +
                    conn->get_outdip().toString() + "]" + ":" + std::to_string(conn->get_dport()) +
                    " rcv_pkt_num:" + std::to_string(rcv_pkt_num) + "\" }";
    } else {
        LOG_INFO("tcp conn end reason:{}. source {}:{} --> destination {}:{}. rcv_pkt_num:{}", close_reason,
                 conn->get_sip().toString(), conn->get_sport(), conn->get_dip().toString(), conn->get_dport(),
                 rcv_pkt_num);

        conn_info = "{ \"tcp_conn_end\" : \"" + conn->get_sip().toString() + ":" + std::to_string(conn->get_sport()) +
                    "-->" + conn->get_dip().toString() + ":" + std::to_string(conn->get_dport()) +
                    " rcv_pkt_num:" + std::to_string(rcv_pkt_num) + "\" }";
    }

    LOG_TRACE("target_file:{} conn_info:{}", g_target_file, conn_info);

    std::ofstream ofs(g_target_file, std::ios_base::app);
    ofs << conn_info << std::endl;
}

// tcp 连接收到数据报文处理
void tcp_conn_event_rcv_data_proc(const tcp_stream_sptr conn, const tcp_stream_data_t *data,
                                  std::shared_ptr<void> usr_conn_data, void *usr_data)
{
    (void)usr_data;
    (void)conn;
    session_sptr s = std::static_pointer_cast<session_t>(usr_conn_data);

    if (data->is_forward) { // 正向报文
        LOG_TRACE("submit forward data. source {}:{} --> destination {}:{}, data len:{}",
                  s->sip.toString(), s->sport, s->dip.toString(), s->dport, data->data_len);
    }
    else { // 反向报文
        LOG_TRACE("submit reversed data. source {}:{} --> destination {}:{}, data len:{}",
                  s->dip.toString(), s->dport, s->sip.toString(), s->sport, data->data_len);
    }

    // !TODO:调用上层业务的接口,将数据提交给上层业务
    if(s->detector){
        LOG_DEBUG("start   detector   len: {}", data->data_len);
        s->detector->dissector(s, data->data, data->data_len, data->is_forward);
    }else{
        LOG_ERROR("detector is null");
    }

    if (s->detectSatus == DATABASE_DETECTION_OK) {
        LOG_TRACE("DATABASE DETECTED ! Type : {}", s->dbType);
    }
}

std::string g_cap_task_name;
pcapmgr_sptr g_pcapmgr = std::make_shared<netpcap_mgr>(); // 抓包管理
IpPacketHandler g_ip_handler;                             // 二三层报文解析及分发
L4PacketHandler g_l4_handler;                             // 四层报文处理

int pcap_proc_init()
{
    if (g_pcapmgr->init(g_logger) != true) {
        LOG_ERROR("pcap manager init failed.");
        return -1;
    }
    return 0;
}

int ip_packet_proc_init()
{
    g_ip_handler.init();
    return 0;
}

int l4_packet_proc_init()
{
    g_l4_handler.init();

    // 如果上层业务不需要报文重组，可以关闭报文重组流程，提高性能
    g_l4_handler.set_tcp_reassemble_enable(false);

    // 注册上层业务回调
    tcp_conn_event_cb_t cb;
    cb.tcp_conn_event_rcv_ctrl_pkt = nullptr;
    cb.tcp_conn_event_start = tcp_conn_event_start_proc;
    cb.tcp_conn_event_end = tcp_conn_event_end_proc;
    cb.tcp_conn_event_rcv_data = tcp_conn_event_rcv_data_proc;
    g_l4_handler.register_tcp_conn_event_cb(cb, nullptr);
    return 0;
}

// 抓包报文处理接口
int pcap_pkt_handler (void *pkt, void *ctx)
{
    (void)ctx;
    pcap_intf_pkt_t *pcap_pkt = (pcap_intf_pkt_t*)pkt;
    LOG_TRACE("Pcap packet, link type:{} packet len:{}", pcap_pkt->link_type, pcap_pkt->frame_len);

    packet_meta_t meta;
    meta.clear();
    meta.pkt_len = uint16_t(pcap_pkt->frame_len);
    meta.timestamp = pcap_pkt->ts;
    meta.link_type = (LinkLayerType)pcap_pkt->link_type;
    // 进行 ip 解析
    uint8_t *out_pkt = nullptr;
    if (g_ip_handler.process_packet(pcap_pkt->frame_data, &out_pkt, &meta) != IP_PKT_OK) {
        return -1;
    }

    if (meta.ip_protocol != 6) { // 非tcp协议 不处理
        return -1;
    }

    // 进行 tcp 重组
    (void)g_l4_handler.process_packet(out_pkt + meta.l4_offset, &meta);
    return 0;
}

void print_statistic()
{
    // 获取统计信息
    pcap_pkt_stat_t cap_stat;
    pcap_stat dev_stat;
    g_pcapmgr->pcap_task_get_stat(g_cap_task_name, cap_stat, dev_stat);
    ip_pkt_stat_t ip_stat = g_ip_handler.get_pkt_stat();
    l4_pkt_stat l4_stat = g_l4_handler.get_pkt_stat();
    tcp_statistic_t tcp_stat = g_l4_handler.get_tcp_stat();
    
    LOG_TRACE(
        "Pcap statistic: rx pkt num:{} rx pkt bytes:{} rx drop num:{} tx pkt num:{} tx pkt bytes:{} tx drop num:{}",
        cap_stat.rx_pkt_nums, cap_stat.rx_pkt_bytes, cap_stat.rx_pkt_drops, cap_stat.tx_pkt_nums, cap_stat.tx_pkt_bytes,
        cap_stat.tx_pkt_drops);
    
    LOG_TRACE("ip statistic: rx pkt num:{} rx pkt bytes:{} tx pkt num:{} tx pkt bytes:{}",
        ip_stat.rcv_pkt_num, ip_stat.rcv_pkt_bytes, ip_stat.snd_pkt_num, ip_stat.snd_pkt_bytes);

    LOG_TRACE("l4 statistic: rx pkt num:{} rx pkt bytes:{} tx pkt num:{} tx pkt bytes:{} cached pkt bytes:{}  error "
             "drop num:{} filter drop num:{}",
             l4_stat.rcv_pkt_num, l4_stat.rcv_pkt_bytes, l4_stat.snd_pkt_num, l4_stat.snd_pkt_bytes,
             l4_stat.cached_pkt_num, l4_stat.drop_pkt_num[L4_PKT_ERR], l4_stat.drop_pkt_num[L4_PKT_FILTER]);

    LOG_INFO("tcp statistic: stream total num:{} syn crate:{} data create:{} rx pkt num:{} rx pkt bytes:{} tx pkt num:{} tx pkt bytes:{}",
        tcp_stat.stream_total_num, tcp_stat.stream_syn_create_num, tcp_stat.stream_data_create_num,
        tcp_stat.pkt_rcv_num, tcp_stat.pkt_rcv_data_bytes, tcp_stat.pkt_submit_num, tcp_stat.pkt_submit_bytes); 
}

// 抓包停止处理接口
void pcap_stop_handler(void *ctx)
{
    (void)ctx;
    g_l4_handler.get_tcp_handler()->close_all_tcp_streams();
    print_statistic();
    LOG_INFO("pcap stop.");
}

// 抓包文件处理
int pcapfile_process(std::string &pcap_file)
{
    LOG_INFO("begin to process pcap file {}", pcap_file);
    
    cap_task_t cap_tasker;
    cap_tasker.task_info.task_name = g_cap_task_name;
    cap_tasker.fn_pf_handler = pcap_pkt_handler;
    cap_tasker.fn_pf_stoper = pcap_stop_handler;
    cap_tasker.fn_ctx = nullptr;
    cap_tasker.file_info.dir = "";
    cap_tasker.file_info.file_name = pcap_file;

    bool ret = g_pcapmgr->pcap_task_start_replay(cap_tasker);
    if (ret == false) {
        LOG_ERROR("file {} replay failed", pcap_file);
    }

    LOG_INFO("process pcap file {} done", pcap_file);
    return 0;
}

// 系统初始化
int init()
{
    if (log_init(g_logger) != 0) {
        std::cout << "log_init error!" << std::endl;
        return -1;
    }
    
    if (tcpip_init(g_logger) != 0) {
        LOG_ERROR("tcpip_init error!");
        return -1;
    }

     if (secsmart_protocol::protocol_init(g_logger) != 0) {
        LOG_ERROR("tcpip_init error!");
        return -1;
    }

    if (pcap_proc_init() != 0) {
        LOG_ERROR("pcap_proc_init error!");
        return -1;
    }

    if (ip_packet_proc_init() != 0) {
        LOG_ERROR("ip_packet_proc_init error!");
        return -1;
    }

    if (l4_packet_proc_init() != 0) {
        LOG_ERROR("l4_packet_proc_init error");
        return -1;
    }

    return 0;
}

int main(int argc, char** argv)
{
    /* if (getenv("SNIFFERDB_USELOG"))
    {
        g_withlog = strcmp("true", getenv("SNIFFERDB_USELOG"))==0;
    } */
    if (argc == 2 && std::strcmp(argv[1], "version") == 0)
    {
        sniffer_db_app_print_version();
        return 0;
    }

    if (init() != 0)
    {
        std::cout << "init error" << std::endl;
        return -1;
    }

    LOG_INFO("init done");

    // 获取抓包文件和目标文件，需要完整路径
    std::string pcap_file, target_file;
    if (argc < 3)
    {
        LOG_ERROR("input parameter error!");
        LOG_ERROR("Usage : sniffer_db_app version (print build infos)");
        LOG_ERROR("Usage : sniffer_db_app pcapfile targetfile [maxHitCount (default 5)] [L (print Debug infos)]");
        return -1;
    }
    pcap_file = argv[1];

    target_file = argv[2];
    g_target_file = target_file;
    LOG_INFO("the pcap file is {}, the target file is {}", pcap_file, target_file);
    g_maxThreshold = (argc >= 4) ? atoi(argv[3]) : MYSQL_MATCH_THRESHOLD;

    if (argc == 5 && *argv[4] == 'L')
    {
        tcpip_set_log_level(level::logger_level::debug, true);
        g_maxThreshold = atoi(argv[3]);
    }

    g_cap_task_name = "replay " + pcap_file;

    // 抓包文件处理
    pcapfile_process(pcap_file);

    LOG_INFO("deal done!");
    return 0;
}